{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['001.jpg',\n",
       " '002.jpg',\n",
       " '003.jpg',\n",
       " '004.jpg',\n",
       " '005.jpg',\n",
       " '006.jpg',\n",
       " '007.jpg',\n",
       " '008.jpg',\n",
       " '009.jpg',\n",
       " '010.jpg',\n",
       " '011.jpg',\n",
       " '012.jpg',\n",
       " '013.jpg',\n",
       " '014.jpg',\n",
       " '015.jpg',\n",
       " '016.jpg',\n",
       " '017.jpg',\n",
       " '018.jpg',\n",
       " '019.jpg',\n",
       " '020.jpg',\n",
       " '021.jpg',\n",
       " '022.jpg',\n",
       " '023.jpg',\n",
       " '024.jpg',\n",
       " '025.jpg',\n",
       " '026.jpg',\n",
       " '027.jpg',\n",
       " '028.jpg',\n",
       " '029.jpg',\n",
       " '030.jpg',\n",
       " '031.jpg',\n",
       " '032.jpg',\n",
       " '033.jpg',\n",
       " '034.jpg',\n",
       " '035.jpg',\n",
       " '036.jpg',\n",
       " '037.jpg',\n",
       " '038.jpg',\n",
       " '039.jpg',\n",
       " '040.jpg',\n",
       " '041.jpg',\n",
       " '042.jpg',\n",
       " '043.jpg',\n",
       " '044.jpg',\n",
       " '045.jpg',\n",
       " '046.jpg',\n",
       " '047.jpg',\n",
       " '048.jpg',\n",
       " '049.jpg',\n",
       " '050.jpg',\n",
       " '051.jpg',\n",
       " '052.jpg',\n",
       " '053.jpg',\n",
       " '054.jpg',\n",
       " '055.jpg',\n",
       " '056.jpg',\n",
       " '057.jpg',\n",
       " '058.jpg',\n",
       " '059.jpg',\n",
       " '060.jpg',\n",
       " '061.jpg',\n",
       " '062.jpg',\n",
       " '063.jpg',\n",
       " '064.jpg',\n",
       " '065.jpg',\n",
       " '066.jpg',\n",
       " '067.jpg',\n",
       " '068.jpg',\n",
       " '069.jpg',\n",
       " '070.jpg',\n",
       " '071.jpg',\n",
       " '072.jpg',\n",
       " '073.jpg',\n",
       " '074.jpg',\n",
       " '075.jpg',\n",
       " '076.jpg',\n",
       " '077.jpg',\n",
       " '078.jpg',\n",
       " '079.jpg',\n",
       " '080.jpg',\n",
       " '081.jpg',\n",
       " '082.jpg',\n",
       " '083.jpg',\n",
       " '084.jpg',\n",
       " '085.jpg',\n",
       " '086.jpg',\n",
       " '087.jpg',\n",
       " '088.jpg',\n",
       " '089.jpg',\n",
       " '090.jpg',\n",
       " '091.jpg',\n",
       " '092.jpg',\n",
       " '093.jpg',\n",
       " '094.jpg',\n",
       " '095.jpg',\n",
       " '096.jpg',\n",
       " '097.jpg',\n",
       " '098.jpg',\n",
       " '099.jpg',\n",
       " '100.jpg']"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "from os.path import join, split, isfile\n",
    "\n",
    "image_dir = 'D:\\\\aic19\\\\aic19-track3-test-data\\\\aic19-track3-segmentation\\\\aic19-track3-frames'\n",
    "\n",
    "only_images = [f for f in os.listdir(image_dir) if isfile(join(image_dir, f)) and f[-4:] == '.jpg']\n",
    "only_images.sort()\n",
    "only_images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# k-means only\n",
    "\n",
    "import cv2\n",
    "import numpy as np\n",
    "\n",
    "output_dir = 'D:\\\\aic19\\\\aic19-track3-test-data\\\\aic19-track3-segmentation\\\\aic19-track3-frames-kmeans'\n",
    "os.makedirs(output_dir)\n",
    "\n",
    "for f in only_images:\n",
    "    print(f)\n",
    "    img = cv2.imread(join(image_dir, f))\n",
    "    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)\n",
    "    \n",
    "    Z = hsv.reshape((-1,3))\n",
    "    # convert to np.float32\n",
    "    Z = np.float32(Z)\n",
    "    # define criteria, number of clusters(K) and apply kmeans()\n",
    "    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)\n",
    "    K = 4\n",
    "    ret,label,center=cv2.kmeans(Z,K,None,criteria,10,cv2.KMEANS_RANDOM_CENTERS)\n",
    "    # Now convert back into uint8, and make original image\n",
    "    center = np.uint8(center)\n",
    "    res = center[label.flatten()]\n",
    "    res2 = res.reshape((img.shape))\n",
    "    img_new = cv2.cvtColor(res2, cv2.COLOR_HSV2BGR)\n",
    "    cv2.imwrite(join(output_dir, f[:-4]+'-kmeans.jpg'), img_new)\n",
    "    \n",
    "    # define range of grass in HSV # H in opencv is in (0, 180)\n",
    "    lower = np.array([25,36,32])\n",
    "    upper = np.array([90,255,255])\n",
    "\n",
    "    # Threshold the HSV image to get grass\n",
    "    mask = cv2.inRange(res2, lower, upper)\n",
    "    \n",
    "    cv2.imwrite(join(output_dir, f[:-4]+'-grass-mask.jpg'), mask)\n",
    "    \n",
    "    mask = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR)\n",
    "    img_new = img/2 + mask/2\n",
    "    cv2.imwrite(join(output_dir, f[:-4]+'-grass-img.jpg'), img_new)\n",
    "    \n",
    "    cv2.imshow('img_new', img_new)\n",
    "    cv2.waitKey(1000)\n",
    "cv2.destroyAllWindows()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "count grass: 79751\n",
      "count road: 274073\n",
      "labels.shape = (353824,)\n",
      "trainData.shape = (353824, 3)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 102,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# train SVM\n",
    "\n",
    "import cv2\n",
    "import numpy as np\n",
    "\n",
    "labels = []\n",
    "trainData = []\n",
    "\n",
    "# road\n",
    "img_road = cv2.imread('D:\\\\aic19\\\\aic19-track3-test-data\\\\aic19-track3-segmentation\\\\road.png')\n",
    "img_road = cv2.cvtColor(img_road, cv2.COLOR_BGR2HSV)\n",
    "h, w, c = img_road.shape\n",
    "count_road = 0\n",
    "for i in range(img_road.shape[0]):\n",
    "    for j in range(img_road.shape[1]):\n",
    "        #if sum(img_road[i,j]) < 255*3:\n",
    "        if img_road[i,j,2] < 255:\n",
    "            labels.append(-1)\n",
    "            trainData.append(img_road[i,j])\n",
    "            count_road += 1\n",
    "\n",
    "# grass\n",
    "img_grass = cv2.imread('D:\\\\aic19\\\\aic19-track3-test-data\\\\aic19-track3-segmentation\\\\grass.png')\n",
    "img_grass = cv2.cvtColor(img_grass, cv2.COLOR_BGR2HSV)\n",
    "h, w, c = img_grass.shape\n",
    "count_grass = 0\n",
    "for i in range(img_grass.shape[0]):\n",
    "    for j in range(img_grass.shape[1]):\n",
    "        #if sum(img_grass[i,j]) < 255*3:\n",
    "        if img_grass[i,j,2] < 255:\n",
    "            labels.append(1)\n",
    "            trainData.append(img_grass[i,j])\n",
    "            count_grass += 1\n",
    "\n",
    "print('count grass: %d' % count_grass)\n",
    "print('count road: %d' % count_road)\n",
    "labels= np.array(labels, dtype=np.int32)\n",
    "trainData = np.array(trainData, dtype=np.float32)\n",
    "print('labels.shape = ' + str(labels.shape))\n",
    "print('trainData.shape = ' + str(trainData.shape))\n",
    "            \n",
    "# train SVM\n",
    "svm = cv2.ml.SVM_create()\n",
    "svm.setType(cv2.ml.SVM_C_SVC)\n",
    "svm.setKernel(cv2.ml.SVM_LINEAR)\n",
    "#svm.setKernel(cv2.ml.SVM_RBF)\n",
    "svm.setTermCriteria((cv2.TERM_CRITERIA_MAX_ITER, 100, 1e-6))\n",
    "svm.train(trainData, cv2.ml.ROW_SAMPLE, labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "001.jpg\n",
      "002.jpg\n",
      "003.jpg\n",
      "004.jpg\n",
      "005.jpg\n",
      "006.jpg\n",
      "007.jpg\n",
      "008.jpg\n",
      "009.jpg\n",
      "010.jpg\n",
      "011.jpg\n",
      "012.jpg\n",
      "013.jpg\n",
      "014.jpg\n",
      "015.jpg\n",
      "016.jpg\n",
      "017.jpg\n",
      "018.jpg\n",
      "019.jpg\n",
      "020.jpg\n",
      "021.jpg\n",
      "022.jpg\n",
      "023.jpg\n",
      "024.jpg\n",
      "025.jpg\n",
      "026.jpg\n",
      "027.jpg\n",
      "028.jpg\n",
      "029.jpg\n",
      "030.jpg\n",
      "031.jpg\n",
      "032.jpg\n",
      "033.jpg\n",
      "034.jpg\n",
      "035.jpg\n",
      "036.jpg\n",
      "037.jpg\n",
      "038.jpg\n",
      "039.jpg\n",
      "040.jpg\n",
      "041.jpg\n",
      "042.jpg\n",
      "043.jpg\n",
      "044.jpg\n",
      "045.jpg\n",
      "046.jpg\n",
      "047.jpg\n",
      "048.jpg\n",
      "049.jpg\n",
      "050.jpg\n",
      "051.jpg\n",
      "052.jpg\n",
      "053.jpg\n",
      "054.jpg\n",
      "055.jpg\n",
      "056.jpg\n",
      "057.jpg\n",
      "058.jpg\n",
      "059.jpg\n",
      "060.jpg\n",
      "061.jpg\n",
      "062.jpg\n",
      "063.jpg\n",
      "064.jpg\n",
      "065.jpg\n",
      "066.jpg\n",
      "067.jpg\n",
      "068.jpg\n",
      "069.jpg\n",
      "070.jpg\n",
      "071.jpg\n",
      "072.jpg\n",
      "073.jpg\n",
      "074.jpg\n",
      "075.jpg\n",
      "076.jpg\n",
      "077.jpg\n",
      "078.jpg\n",
      "079.jpg\n",
      "080.jpg\n",
      "081.jpg\n",
      "082.jpg\n",
      "083.jpg\n",
      "084.jpg\n",
      "085.jpg\n",
      "086.jpg\n",
      "087.jpg\n",
      "088.jpg\n",
      "089.jpg\n",
      "090.jpg\n",
      "091.jpg\n",
      "092.jpg\n",
      "093.jpg\n",
      "094.jpg\n",
      "095.jpg\n",
      "096.jpg\n",
      "097.jpg\n",
      "098.jpg\n",
      "099.jpg\n",
      "100.jpg\n"
     ]
    }
   ],
   "source": [
    "# predict by SVM (kmeans)\n",
    "\n",
    "import os\n",
    "\n",
    "output_dir = 'D:\\\\aic19\\\\aic19-track3-test-data\\\\aic19-track3-segmentation\\\\aic19-track3-frames-svm-k4'\n",
    "if not os.path.isdir(output_dir):\n",
    "    os.makedirs(output_dir)\n",
    "\n",
    "for f in only_images:\n",
    "    print(f)\n",
    "    img = cv2.imread(join(image_dir, f))\n",
    "    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)\n",
    "    \n",
    "    Z = hsv.reshape((-1,3))\n",
    "    # convert to np.float32\n",
    "    Z = np.float32(Z)\n",
    "    # define criteria, number of clusters(K) and apply kmeans()\n",
    "    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)\n",
    "    K = 4\n",
    "    ret,label,center=cv2.kmeans(Z,K,None,criteria,10,cv2.KMEANS_RANDOM_CENTERS)\n",
    "    # Now convert back into uint8, and make original image\n",
    "    center = np.uint8(center)\n",
    "    res = center[label.flatten()]\n",
    "    res2 = res.reshape((img.shape))\n",
    "    \n",
    "    # classify using SVM\n",
    "    sample = res2.reshape(-1, 3).astype(np.float32)\n",
    "    \n",
    "    _, response = svm.predict(sample)\n",
    "    response = response.reshape(img.shape[0], img.shape[1])\n",
    "    \n",
    "    indices_grass = np.argwhere(response < 0)\n",
    "    indices_road = np.argwhere(response > 0)\n",
    "    \n",
    "    mask_grass = np.zeros((img.shape[0], img.shape[1]), dtype=np.uint8)\n",
    "    mask_grass[indices_grass[:,0] ,indices_grass[:,1]] = 255\n",
    "    \n",
    "    mask_road = np.zeros((img.shape[0], img.shape[1]), dtype=np.uint8)\n",
    "    mask_road[indices_road[:,0] ,indices_road[:,1]] = 255\n",
    "    \n",
    "    \n",
    "    res2 = cv2.cvtColor(res2, cv2.COLOR_HSV2BGR)\n",
    "    cv2.imwrite(join(output_dir, f[:-4]+'-kmeans.jpg'), res2)\n",
    "    cv2.imwrite(join(output_dir, f[:-4]+'-grass-mask.png'), mask_grass)\n",
    "    cv2.imwrite(join(output_dir, f[:-4]+'-road-mask.png'), mask_road)\n",
    "    \n",
    "    mask_grass = cv2.cvtColor(mask_grass, cv2.COLOR_GRAY2BGR)\n",
    "    img_grass = img/2 + mask_grass/2\n",
    "    cv2.imwrite(join(output_dir, f[:-4]+'-grass-img.jpg'), img_grass)\n",
    "\n",
    "    mask_road = cv2.cvtColor(mask_road, cv2.COLOR_GRAY2BGR)\n",
    "    img_road = img/2 + mask_road/2\n",
    "    cv2.imwrite(join(output_dir, f[:-4]+'-road-img.jpg'), img_road)\n",
    "\n",
    "    #cv2.imshow('img_grass', img_grass)\n",
    "    #cv2.waitKey(1000)\n",
    "#cv2.destroyAllWindows()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "001.jpg\n",
      "002.jpg\n",
      "003.jpg\n",
      "004.jpg\n",
      "005.jpg\n",
      "006.jpg\n",
      "007.jpg\n",
      "008.jpg\n",
      "009.jpg\n",
      "010.jpg\n",
      "011.jpg\n",
      "012.jpg\n",
      "013.jpg\n",
      "014.jpg\n",
      "015.jpg\n",
      "016.jpg\n",
      "017.jpg\n",
      "018.jpg\n",
      "019.jpg\n",
      "020.jpg\n",
      "021.jpg\n",
      "022.jpg\n",
      "023.jpg\n",
      "024.jpg\n",
      "025.jpg\n",
      "026.jpg\n",
      "027.jpg\n",
      "028.jpg\n",
      "029.jpg\n",
      "030.jpg\n",
      "031.jpg\n",
      "032.jpg\n",
      "033.jpg\n",
      "034.jpg\n",
      "035.jpg\n",
      "036.jpg\n",
      "037.jpg\n",
      "038.jpg\n",
      "039.jpg\n",
      "040.jpg\n",
      "041.jpg\n",
      "042.jpg\n",
      "043.jpg\n",
      "044.jpg\n",
      "045.jpg\n",
      "046.jpg\n",
      "047.jpg\n",
      "048.jpg\n",
      "049.jpg\n",
      "050.jpg\n",
      "051.jpg\n",
      "052.jpg\n",
      "053.jpg\n",
      "054.jpg\n",
      "055.jpg\n",
      "056.jpg\n",
      "057.jpg\n",
      "058.jpg\n",
      "059.jpg\n",
      "060.jpg\n",
      "061.jpg\n",
      "062.jpg\n",
      "063.jpg\n",
      "064.jpg\n",
      "065.jpg\n",
      "066.jpg\n",
      "067.jpg\n",
      "068.jpg\n",
      "069.jpg\n",
      "070.jpg\n",
      "071.jpg\n",
      "072.jpg\n",
      "073.jpg\n",
      "074.jpg\n",
      "075.jpg\n",
      "076.jpg\n",
      "077.jpg\n",
      "078.jpg\n",
      "079.jpg\n",
      "080.jpg\n",
      "081.jpg\n",
      "082.jpg\n",
      "083.jpg\n",
      "084.jpg\n",
      "085.jpg\n",
      "086.jpg\n",
      "087.jpg\n",
      "088.jpg\n",
      "089.jpg\n",
      "090.jpg\n",
      "091.jpg\n",
      "092.jpg\n",
      "093.jpg\n",
      "094.jpg\n",
      "095.jpg\n",
      "096.jpg\n",
      "097.jpg\n",
      "098.jpg\n",
      "099.jpg\n",
      "100.jpg\n"
     ]
    }
   ],
   "source": [
    "# predict by SVM (every pixel)\n",
    "\n",
    "import os\n",
    "\n",
    "output_dir = 'D:\\\\aic19\\\\aic19-track3-test-data\\\\aic19-track3-segmentation\\\\aic19-track3-frames-svm'\n",
    "if not os.path.isdir(output_dir):\n",
    "    os.makedirs(output_dir)\n",
    "\n",
    "for f in only_images:\n",
    "    print(f)\n",
    "    img = cv2.imread(join(image_dir, f))\n",
    "    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)\n",
    "    \n",
    "    # classify using SVM\n",
    "    sample = hsv.reshape(-1, 3).astype(np.float32)\n",
    "    \n",
    "    _, response = svm.predict(sample)\n",
    "    response = response.reshape(img.shape[0], img.shape[1])\n",
    "    \n",
    "    indices_grass = np.argwhere(response < 0)\n",
    "    indices_road = np.argwhere(response > 0)\n",
    "    \n",
    "    mask_grass = np.zeros((img.shape[0], img.shape[1]), dtype=np.uint8)\n",
    "    mask_grass[indices_grass[:,0] ,indices_grass[:,1]] = 255\n",
    "    \n",
    "    mask_road = np.zeros((img.shape[0], img.shape[1]), dtype=np.uint8)\n",
    "    mask_road[indices_road[:,0] ,indices_road[:,1]] = 255\n",
    "    \n",
    "    cv2.imwrite(join(output_dir, f[:-4]+'.jpg'), img)\n",
    "    cv2.imwrite(join(output_dir, f[:-4]+'-grass-mask.png'), mask_grass)\n",
    "    cv2.imwrite(join(output_dir, f[:-4]+'-road-mask.png'), mask_road)\n",
    "    \n",
    "    mask_grass = cv2.cvtColor(mask_grass, cv2.COLOR_GRAY2BGR)\n",
    "    img_grass = img/2 + mask_grass/2\n",
    "    cv2.imwrite(join(output_dir, f[:-4]+'-grass-img.jpg'), img_grass)\n",
    "\n",
    "    mask_road = cv2.cvtColor(mask_road, cv2.COLOR_GRAY2BGR)\n",
    "    img_road = img/2 + mask_road/2\n",
    "    cv2.imwrite(join(output_dir, f[:-4]+'-road-img.jpg'), img_road)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
